# -*- coding: utf-8 -*-

import os
import datetime
import pandas as pd
import logging

from pyquant import SignalEvent, Strategy, CSVDataHandler, SimulatedExecutionHandler, BasicPortfolio, Backtest, MarketOrder, LimitOrder


class HelloworldStrategy(Strategy):

    def init(self, context):
        """
        init 必须实现,
        必须初始化stocks数组, 初始化之后引擎会根据数组发送bar数据.
        """
        context.stocks = ['600008', '600018']
        context.long_window = 10
        context.short_window = 5
        context.bought = {}
        context.bought['600008'] = False
        context.bought['600018'] = False



    def handle_bar(self, context, bar_dict):
        """
        仅测试, 无任何意义
        """
        if self.event.type == 'BAR':
            for s in context.stocks:
                bar = context.data_handler.get_latest_bar(s)
                if bar is None or bar == []: continue
                bars = context.data_handler.get_latest_bars(s, N=context.long_window)
                if len(bars) >= context.long_window:
                    df = pd.DataFrame(bars, columns=['symbol', 'datetime', 'open', 'high', 'low', 'close', 'volume'])
                    df['MA_l'] = df['close'].rolling(context.long_window, min_periods=1).mean()
                    df['MA_s'] = df['close'].rolling(context.short_window, min_periods=1).mean()
                    if df['MA_l'].iloc[-1] < df['MA_s'].iloc[-1] and df['MA_l'].iloc[-2] > df['MA_s'].iloc[-2]:
                        if not context.bought[s]:
                            # signal = SignalEvent(bar.symbol, bar.datetime, 'LONG')
                            # context.events.put(signal)
                            # self.order_shares(s, 8800)
                            # self.order_value(s, 10000)
                            self.order_target_percent(s, 0.5)
                            context.bought[s] = True
                    elif df['MA_l'].iloc[-1] < df['MA_s'].iloc[-1] and df['MA_l'].iloc[-2] < df['MA_s'].iloc[-2]:
                        if context.bought[s]:
                            pass
                            # signal = SignalEvent(bar.symbol, bar.datetime, 'EXIT')
                            # context.events.put(signal)
                            # self.order_value(s, -10000)
                            # self.order_target_percent(s, 0)
                            # self.order_shares(s, -8800)
                            # context.bought[s] = False



if __name__ == '__main__':
    csv_dir = os.path.join(os.getcwd(), 'testdata')
    # print(csv_dir)
    symbol_list = ['600008', '600018']
    initial_capital = 100000.0
    heartbeat = 0.0
    start_date = datetime.datetime(2015, 11, 2, 0, 0)
    end_date = datetime.datetime(2015, 12, 30, 23, 59)

    backtest = Backtest(initial_capital,
                        SimulatedExecutionHandler,
                        BasicPortfolio, HelloworldStrategy,
                        slippage_type='fixed', commission_type='default', csv_dir=csv_dir, start_date=start_date, end_date=end_date)

    backtest.setup_logger(logging.DEBUG)
    positions, holdings = backtest.run_backtest()
    print(holdings.tail())
    backtest.plot()
